﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Text;
using System.Windows.Forms;
using System.Diagnostics;
using System.IO;

namespace iMan
{
    public partial class MainForm : Form
    {
        protected const int CMD_REFRESH = 1;
        protected IpHelper helper;
        protected Dictionary<string, MIB_TCPROW_OWNER_PID> innerExtendedTcpRowList;
        protected Dictionary<string, MIB_UDPROW_OWNER_PID> innerExtendedUdpRowList;
        protected Dictionary<string, MIB_EXTCPROW> innerExTcpRowList;
        protected Dictionary<string, MIB_EXUDPROW> innerExUdpRowList;
        protected int currentPID;
        protected string currentFolder;

        public MainForm()
        {
            InitializeComponent();
            Load += new EventHandler(MainForm_Load);
            FormClosing += new FormClosingEventHandler(MainForm_FormClosing);

            innerExtendedTcpRowList = new Dictionary<string, MIB_TCPROW_OWNER_PID>();
            innerExtendedUdpRowList = new Dictionary<string, MIB_UDPROW_OWNER_PID>();
            innerExTcpRowList = new Dictionary<string, MIB_EXTCPROW>();
            innerExUdpRowList = new Dictionary<string, MIB_EXUDPROW>();
            helper = new IpHelper();

            currentFolder = null;
            currentPID = -1;

            if (Environment.OSVersion.Version.Major <= 5)
            {
                helper.GetExTcpTable();
                helper.GetExUdpTable();
            }
            else
            {
                helper.GetExtendedTcpTable();
                helper.GetExtendedUdpTable();
            }
        }

        private void MainForm_FormClosing(object sender, FormClosingEventArgs e)
        {
            WMIVendor.Instance.Dispose();
        }

        private void MainForm_Load(object sender, EventArgs e)
        {
            Text += " - [" + Application.ProductVersion + "]";
            IntPtr systemMenu = User32.GetSystemMenu(base.Handle, false);
            User32.InsertMenu(systemMenu, 7, MenuFlag.MF_SEPARATOR | MenuFlag.MF_BYPOSITION, (uint)SysCommand.SC_SEPARATOR, "");
            User32.InsertMenu(systemMenu, 8, MenuFlag.MF_BYPOSITION, CMD_REFRESH, "Refresh");

            if (Environment.OSVersion.Version.Major <= 5)
            {
                HandleExTable();
            }
            else
            {
                HandleExtendedTable();
            }
        }

        protected override void WndProc(ref Message m)
        {
            if ((m.Msg == (int)Msgs.WM_SYSCOMMAND) && (((int)m.WParam) == CMD_REFRESH))
            {
                ReloadTables();
            }
            base.WndProc(ref m);
        }

        private void ReloadTables()
        {
            TreeNode root = treeView1.Nodes["tn_Root"];
            if (root != null)
            {
                TreeNode UDPNode = root.Nodes["tn_UDP"];
                TreeNode TCPNode = root.Nodes["tn_TCP"];
                if (UDPNode != null)
                    UDPNode.Nodes.Clear();
                if (TCPNode != null)
                    TCPNode.Nodes.Clear();
            }
            propertyGrid1.SelectedObject = null;
            listView1.Items.Clear();
            innerExtendedTcpRowList = new Dictionary<string, MIB_TCPROW_OWNER_PID>();
            innerExtendedUdpRowList = new Dictionary<string, MIB_UDPROW_OWNER_PID>();
            innerExTcpRowList = new Dictionary<string, MIB_EXTCPROW>();
            innerExUdpRowList = new Dictionary<string, MIB_EXUDPROW>();
            helper = new IpHelper();

            if (Environment.OSVersion.Version.Major <= 5)
            {
                helper.GetExTcpTable();
                helper.GetExUdpTable();
                HandleExTable();
            }
            else
            {
                helper.GetExtendedTcpTable();
                helper.GetExtendedUdpTable();
                HandleExtendedTable();
            }
        }

        private void HandleExTable()
        {
            try
            {
                foreach (MIB_EXTCPROW item in helper.TcpExTable.table)
                {
                    int port = item.Local.Port;
                    string name = item.ProcessName;
                    string key = port.ToString() + name;
                    innerExTcpRowList[key] = item;
                    AppendTreeNode(port, name, true);
                    ListViewItem vItem = new ListViewItem(new string[] { item.Local.ToString(), item.Remote.ToString(), "TCP", item.dwProcessId.ToString(), item.ProcessName, item.StrgState });
                    listView1.Items.Add(vItem);
                }
                foreach (MIB_EXUDPROW item in helper.UdpExTable.table)
                {
                    int port = item.Local.Port;
                    string name = item.ProcessName;
                    string key = port.ToString() + name;
                    innerExUdpRowList[key] = item;
                    AppendTreeNode(port, name, false);
                    ListViewItem vItem = new ListViewItem(new string[] { item.Local.ToString(), "", "UDP", item.dwProcessId.ToString(), item.ProcessName, "" });
                    listView1.Items.Add(vItem);
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex.Message);
            }
        }

        private void HandleExtendedTable()
        {
            try
            {
                foreach (MIB_TCPROW_OWNER_PID item in helper.TcpExtendedTable.table)
                {
                    int port = item.Local.Port;
                    string name = item.ProcessName;
                    string key = port.ToString() + name;
                    innerExtendedTcpRowList[key] = item;
                    AppendTreeNode(port, name, true);
                    ListViewItem vItem = new ListViewItem(new string[] { item.Local.ToString(), item.Remote.ToString(), "TCP", item.dwOwningPid.ToString(), item.ProcessName, item.State });
                    listView1.Items.Add(vItem);
                }

                foreach (MIB_UDPROW_OWNER_PID item in helper.UdpExtendedTable.table)
                {
                    int port = item.Local.Port;
                    string name = item.ProcessName;
                    string key = port.ToString() + name;
                    innerExtendedUdpRowList[key] = item;
                    AppendTreeNode(port, name, false);
                    ListViewItem vItem = new ListViewItem(new string[] { item.Local.ToString(), "", "UDP", item.dwOwningPid.ToString(), item.ProcessName, "" });
                    listView1.Items.Add(vItem);
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex.Message);
            }
        }

        private void AppendTreeNode(int port, string name, bool isTcp)
        {
            TreeNode root = treeView1.Nodes["tn_Root"];
            if (root == null) return;
            TreeNode protocolNode = root.Nodes["tn_UDP"];
            if (isTcp)
            {
                protocolNode = root.Nodes["tn_TCP"];
            }
            if (protocolNode == null) return;
            TreeNode portNode = protocolNode.Nodes[port.ToString()];
            if (portNode == null)
            {
                protocolNode.Nodes.Add(port.ToString(), port.ToString());
                portNode = protocolNode.Nodes[port.ToString()];
            }
            TreeNode nameNode = portNode.Nodes[name];
            if (nameNode == null) portNode.Nodes.Add(name, name);
        }

        private void treeView1_AfterSelect(object sender, TreeViewEventArgs e)
        {
            if (e.Node == null) return;
            TreeNode nameNode = e.Node;
            TreeNode portNode = nameNode.Parent;
            if (portNode == null) return;
            TreeNode protocolNode = portNode.Parent;
            if (portNode != null && protocolNode != null)
            {
                string key = portNode.Name + nameNode.Name;
                bool isTcp = true;
                if (protocolNode.Name == "tn_UDP") isTcp = false;
                OnItemSelected(key, isTcp);
            }
            else propertyGrid1.SelectedObject = null;
        }

        private void OnItemSelected(string key, bool isTcp)
        {
            if (key == null || key.Trim() == "") return;

            if (isTcp)
            {
                if (innerExtendedTcpRowList.ContainsKey(key))
                {
                    MIB_TCPROW_OWNER_PID item = innerExtendedTcpRowList[key];
                    UpdateStatus(string.Format("Retrieving process [{0}] information...", item.ProcessName));
                    WMIProcessInfo info = WMIVendor.Instance.GetProcessInfoBy(item.dwOwningPid);
                    currentPID = item.dwOwningPid;
                    currentFolder = Path.GetDirectoryName(info.ExecutablePath);
                    UpdateStatus("Ready");
                    propertyGrid1.SelectedObject = info;
                }
                else if (innerExTcpRowList.ContainsKey(key))
                {
                    MIB_EXTCPROW item = innerExTcpRowList[key];
                    UpdateStatus(string.Format("Retrieving process [{0}] information...", item.ProcessName));
                    WMIProcessInfo info = WMIVendor.Instance.GetProcessInfoBy(item.dwProcessId);
                    currentPID = item.dwProcessId;
                    currentFolder = Path.GetDirectoryName(info.ExecutablePath);
                    UpdateStatus("Ready");
                    propertyGrid1.SelectedObject = info;
                }
                else
                {
                    propertyGrid1.SelectedObject = null;
                    currentPID = -1;
                    currentFolder = null;
                }
            }
            else
            {
                if (innerExtendedUdpRowList.ContainsKey(key))
                {
                    MIB_UDPROW_OWNER_PID item = innerExtendedUdpRowList[key];
                    UpdateStatus(string.Format("Retrieving process [{0}] information...", item.ProcessName));
                    WMIProcessInfo info = WMIVendor.Instance.GetProcessInfoBy(item.dwOwningPid);
                    currentPID = item.dwOwningPid;
                    currentFolder = Path.GetDirectoryName(info.ExecutablePath);
                    UpdateStatus("Ready");
                    propertyGrid1.SelectedObject = info;
                }
                else if (innerExUdpRowList.ContainsKey(key))
                {
                    MIB_EXUDPROW item = innerExUdpRowList[key];
                    UpdateStatus(string.Format("Retrieving process [{0}] information...", item.ProcessName));
                    WMIProcessInfo info = WMIVendor.Instance.GetProcessInfoBy(item.dwProcessId);
                    currentPID = item.dwProcessId;
                    currentFolder = Path.GetDirectoryName(info.ExecutablePath);
                    UpdateStatus("Ready");
                    propertyGrid1.SelectedObject = info;
                }
                else
                {
                    propertyGrid1.SelectedObject = null;
                    currentPID = -1;
                    currentFolder = null;
                }
            }
        }

        private void listView1_Click(object sender, EventArgs e)
        {
            if (listView1.SelectedItems == null || listView1.SelectedItems.Count == 0) return;
            ListViewItem item = listView1.SelectedItems[0];
            string localEP = item.SubItems[0].Text;
            if (!localEP.Contains(":")) return;
            string port = localEP.Split(':')[1];
            string protocol = item.SubItems[2].Text;
            int.TryParse(item.SubItems[3].Text, out currentPID);
            string name = item.SubItems[4].Text;
            OnItemSelected(port + name, protocol == "TCP");
        }

        private void UpdateStatus(string statusInfo)
        {
            if (InvokeRequired)
            {
                MethodInvoker mi = delegate
                {
                    lbl_Status.Text = statusInfo;
                };
                Invoke(mi);
            }
            else
            {
                lbl_Status.Text = statusInfo;
            }
        }

        private void ts_KillTask_Click(object sender, EventArgs e)
        {
            if (currentPID > -1)
            {
                try
                {
                    if (MessageBox.Show("Kill the selected Task?", "iMan", MessageBoxButtons.YesNo, MessageBoxIcon.Information) == DialogResult.Yes)
                    {
                        Process.GetProcessById(currentPID).Kill();
                    }
                }
                catch (Exception ex)
                {
                    Console.WriteLine(ex.Message);
                }
                finally
                {
                    ReloadTables();
                }
            }
        }

        private void ts_OpenFolder_Click(object sender, EventArgs e)
        {
            if (currentFolder != null)
            {
                if (Directory.Exists(currentFolder))
                {
                    Process.Start(currentFolder);
                }
            }
        }

        private void contextMenuStrip1_Opening(object sender, CancelEventArgs e)
        {
            if (listView1.SelectedItems.Count == 0)
            {
                ts_KillTask.Enabled = false;
                ts_OpenFolder.Enabled = false;
            }
            if (currentFolder == null || currentFolder.Trim() == "")
            {
                ts_OpenFolder.Enabled = false;
            }
        }
    }
}
